# -*- coding: utf-8 -*-

'''
Moduł zawierający obsługę parsowania plików z sekwencjami w formacie FASTA i CLUSTALW.

'''

from Bio import AlignIO
##Szerokość odstępu pomiędzy nazwą sekwencji a sekwencją
NAME_SEQ_GAP_LEN = 6
##Odstęp pomiędzy nazwą sekwencji a sekwencją
NAME_SEQ_GAP = ' ' * NAME_SEQ_GAP_LEN
##Maksymalna liczba nukleotydów w wierszu
SEQ_LEN = 50

##Funkcja odczytująca 3 sekwencje w formacie FASTA z pliku tekstowego.
#@param fp uchwyt na otwarty plik tekstowy
#@return kolejne wywołania zwracają kolejną listę [nazwa, sekwencja]
def read_fasta(fp):
    name, seq = None, []
    if fp:
        for line in fp:
            line = line.rstrip()
            if line.startswith(">"):
                if name: 
                    yield (name[1:], ''.join(seq))
                name, seq = line, []
            else:
                seq.append(line)

    if name: 
        yield (name[1:], ''.join(seq))
    else:
        yield (name, seq)
        
##Funkcja odczytująca 3 dopasowane sekwencje w formacie CLUSTALW z pliku tekstowego.
#@param fp uchwyt na otwarty plik tekstowy
#@return listę nazw i listę sekwencji, None,None w przypadku błędu
def read_clustal(fp):
    nameList = []
    seqList = []
    try:
        c_align = AlignIO.read(fp, "clustal")
    except:
        return None, None
        
    for record in c_align:
        nameList.append(record.id)
        seqList.append(record.seq.tostring())
        
    return nameList, seqList

##Funkcja zapisująca 3 dopasowane sekwencje w formacie CLUSTALW do pliku tekstowego.
#@param fp uchwyt na otwarty plik tekstowy
#@param nameList lista nazw kolejnych sekwencji
#@param seqList lista kolejnych sekwencji
#@param exactList lista zawierająca '*', gdy wszystkie elementy sekwencji takie same lub ' ', gdy różne
def save_clustal(fp, nameList, seqList, exactList):
    maxiName = max(map(len, nameList))
    seqStartPos = maxiName + NAME_SEQ_GAP_LEN
    maxIter = len(seqList[0]) / SEQ_LEN
    
    if fp:
        fp.write('CLUSTAL 2.1 multiple sequence alignment\n\n')
        
        for i in range(maxIter + 1):
            fp.write('\n')
            for j in range(len(nameList)):
                fp.write(''.join(nameList[j]) + ' ' * (maxiName - len(nameList[j])) + NAME_SEQ_GAP + ''.join(seqList[j][i*SEQ_LEN : (i+1)*SEQ_LEN]) + '\n')
            fp.write(' ' * seqStartPos + ''.join(exactList[i*SEQ_LEN : (i+1)*SEQ_LEN]) + '\n')

    else:
        raise TypeError("Incorrect output file passed (None).")
        